#include <set>
#include "Sort.h"

Sort::Sort(int max_age, int min_hits, double iou_threshold) :
max_age_(max_age), min_hits_(min_hits), iou_threshold_(iou_threshold), frame_count_(0)
{
}

vector<TrackingBox> Sort::update(vector<TrackingBox> & dets)
{
	vector<Rect_<float>> predictedBoxes;
	vector<vector<double>> iouMatrix;
	vector<int> assignment;
	set<int> unmatchedDetections;
	set<int> unmatchedTrajectories;
	set<int> allItems;
	set<int> matchedItems;
	vector<cv::Point> matchedPairs;
	vector<TrackingBox> frameTrackingResult;
	unsigned int trkNum = 0;
	unsigned int detNum = 0;

	frame_count_++;

    if (trackers_.size() == 0) // the first frame met
    {
        // initialize kalman trackers_ using first detections.
        for (unsigned int i = 0; i < dets.size(); i++)
        {
            KalmanTracker trk = KalmanTracker(dets[i].box);
            trackers_.push_back(trk);
        }
        // output the first frame detections
        for (unsigned int id = 0; id < dets.size(); id++)
        {
            TrackingBox tb = dets[id];
			tb.id = id + 1;
            // resultsFile << tb.frame << "," << id + 1 << "," << tb.box.x << "," << tb.box.y << "," << tb.box.width << "," << tb.box.height << ",1,-1,-1,-1" << endl;
			frameTrackingResult.push_back(tb);

        }
        // continue;
		return frameTrackingResult;
    }

	// 3.1. get predicted locations from existing trackers_.

	for (auto it = trackers_.begin(); it != trackers_.end();)
	{
		Rect_<float> pBox = (*it).predict();
		if (pBox.x >= 0 && pBox.y >= 0)
		{
			predictedBoxes.push_back(pBox);
			it++;
		}
		else
		{
			it = trackers_.erase(it);
			//cerr << "Box invalid at frame: " << frame_count << endl;
		}
	}

	// 3.2. associate detections to tracked object (both represented as bounding boxes)
	// dets : dets
	trkNum = predictedBoxes.size();
	detNum = dets.size();

	// iouMatrix.clear();
	iouMatrix.resize(trkNum, vector<double>(detNum, 0));

	for (unsigned int i = 0; i < trkNum; i++) // compute iou matrix as a distance matrix
	{
		for (unsigned int j = 0; j < detNum; j++)
		{
			// use 1-iou because the hungarian algorithm computes a minimum-cost assignment.
			iouMatrix[i][j] = 1 - GetIOU(predictedBoxes[i], dets[j].box);
		}
	}

	// solve the assignment problem using hungarian algorithm.
	// the resulting assignment is [track(prediction) : detection], with len=preNum
	HungarianAlgorithm HungAlgo;
	HungAlgo.Solve(iouMatrix, assignment);

	// find matches, unmatched_detections and unmatched_predictions

	if (detNum > trkNum) //	there are unmatched detections
	{
		for (unsigned int n = 0; n < detNum; n++)
			allItems.insert(n);

		for (unsigned int i = 0; i < trkNum; ++i)
			matchedItems.insert(assignment[i]);

		set_difference(allItems.begin(), allItems.end(),
			matchedItems.begin(), matchedItems.end(),
			insert_iterator<set<int>>(unmatchedDetections, unmatchedDetections.begin()));
	}
	else if (detNum < trkNum) // there are unmatched trajectory/predictions
	{
		for (unsigned int i = 0; i < trkNum; ++i)
			if (assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm
				unmatchedTrajectories.insert(i);
	}
	else
		;

	// filter out matched with low IOU
	for (unsigned int i = 0; i < trkNum; ++i)
	{
		if (assignment[i] == -1) // pass over invalid values
			continue;
		if (1 - iouMatrix[i][assignment[i]] < iou_threshold_)
		{
			unmatchedTrajectories.insert(i);
			unmatchedDetections.insert(assignment[i]);
		}
		else
			matchedPairs.push_back(cv::Point(i, assignment[i]));
	}

	// 3.3. updating trackers_

	// update matched trackers_ with assigned detections.
	// each prediction is corresponding to a tracker
	int detIdx, trkIdx;
	for (unsigned int i = 0; i < matchedPairs.size(); i++)
	{
		trkIdx = matchedPairs[i].x;
		detIdx = matchedPairs[i].y;
		trackers_[trkIdx].update(dets[detIdx].box);
	}

	// create and initialise new trackers_ for unmatched detections
	for (auto umd : unmatchedDetections)
	{
		KalmanTracker tracker = KalmanTracker(dets[umd].box);
		trackers_.push_back(tracker);
	}

	// get trackers_' output
	for (auto it = trackers_.begin(); it != trackers_.end();)
	{
		if (((*it).m_time_since_update < 1) &&
			((*it).m_hit_streak >= min_hits_ || frame_count_ <= min_hits_))
		{
			TrackingBox res;
			res.box = (*it).get_state();
			res.id = (*it).m_id + 1;
			res.frame = frame_count_;
			frameTrackingResult.push_back(res);
			it++;
		}
		else
			it++;

		// remove dead tracklet
		if (it != trackers_.end() && (*it).m_time_since_update > max_age_)
			it = trackers_.erase(it);
	}

	return frameTrackingResult;
}